import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
import scipy


class Generate_Data():
    def __init__(self, N_stocks=500, N_features=5, N_T=300):
        self.N_stocks = N_stocks
        self.N_features = N_features
        self.N_T = N_T

    def gen(self):

        rand_specific_return_coef = 0.01
        rand_specific_return = np.random.randn(self.N_stocks, self.N_T) * rand_specific_return_coef

        def gen_rand(t, dim, mean=0):
            import mfm.levy_process as levy_process
            import mfm.levy_measure as levy_measure
            lvy_measure = levy_measure.MertonMeasure(mean=mean, sigma=0.001, lambda_value=0.01)
            lvy = levy_process.MertonProcess(volatility=1.01, rate_free=0, t=10, merton_measure=lvy_measure, origin=0)
            s = {}
            for t in range(self.N_stocks):
                x = [0]
                for _ in range(self.N_T - 1):
                    x.append(x[-1] + lvy.next_jump_size())
                x = pd.Series(x)
                s[t] = x
            s = pd.DataFrame(s)
            return s

        signal = []
        
        for u in range(self.N_features):
            s = gen_rand(self.N_T, self.N_stocks, mean=0.0001)
            from scipy.linalg import eigh, cholesky
            rho_temp = 0.9
            ground_truth_corr_matrix = np.matrix(
                [[1 if i == j else rho_temp for j in range(self.N_stocks)] for i in range(self.N_stocks)])
            ground_truth_signal = cholesky(ground_truth_corr_matrix)

            signal.append(np.einsum('ij,jk->ik', s.values, ground_truth_signal))
        X = np.array(signal)
        f = gen_rand(self.N_T, self.N_features, mean=0.0001).values

        R_noise = np.random.randn(self.N_T, self.N_stocks) * 0.001

        R_simulated = pd.DataFrame(np.einsum('ijk,jk->jk', X, f))
        R = R_simulated.diff()
        R = R.fillna(0)
        R = R + R_noise

        R_cum = R.cumsum()
        return X, R
